(* $Id$ *)

(** The following will be imported from the standard library in the future. *)

(* [&] decides the boolean conjunction. *)
deferred val & : (bool * bool) -> bool
where (forall (x, y). pre (&) (x, y))
  and (forall (x, y). post (&) (x, y) (true) <=> (x = true and y = true))

(* [+'] computes the addition between two integers. *)
deferred val +? : (int * int) -> int
where (forall (x, y). pre (+?) (x, y))
  and (forall (x, y, z). post (+?) (x, y) (z) <=> z = x + y)

(* [-'] computes the substraction between two integers. *)
deferred val -? : (int * int) -> int
where (forall (x, y). pre (-?) (x, y))
  and (forall (x, y, z). post (-?) (x, y) (z) <=> z = x - y)

(* [max] is the logical function that returns the maximum of two integers. *)
deferred logic max  : (int * int) ~> int
where forall (x, y). (x < y => max (x, y) = y) and (y <= x => max (x, y) = x)

(* [max'] is the implementation of the [max] logical function. *)
deferred val max? : (int * int) -> int
where (forall (x, y). pre (max?) (x, y))
  and (forall (x, y, z). post (max?) (x, y) (z) <=> z = max (x, y))

(* The bounded distance between two integers. *)
deferred predicate dist : (int * int * int) ~> prop
where forall (x, y, p). dist (x, y, p) <=> ((x - y >= (0 - p) and (x - y) <= p))

(* Comparison operators over integers. *)
deferred val <=?  : (int * int) -> bool where
    (forall (x, y). (post (<=?) (x, y) (true)  <=> x <= y))
and (forall (x, y). (post (<=?) (x, y) (false) <=> not (x <= y)))
and (forall (x, y). pre (<=?) (x, y))

deferred val >=?  : (int * int) -> bool where
    (forall (x, y). (post (>=?) (x, y) (true)  <=> x >= y))
and (forall (x, y). (post (>=?) (x, y) (false) <=> not (x >= y)))
and (forall (x, y). pre (>=?) (x, y))

deferred val >?   : (int * int) -> bool where
    (forall (x, y). (post (>?) (x, y) (true)  <=> x > y))
and (forall (x, y). (post (>?) (x, y) (false) <=> not (x > y)))
and (forall (x, y). pre (>?) (x, y))

deferred val =?   : (int * int) -> bool
where (forall (x, y). pre ( =? ) (x, y))
  and (forall (x, y). post (=?) (x, y) (true) <=> x = y)
  and (forall (x, y). post (=?) (x, y) (false) <=> not (x = y))

(** Elements of the sets. *)

(* The type of the elements. *)
deferred type elt : *

(* Their equality is decidable. *)
deferred val =elt?  : (elt * elt) -> bool
where (forall x. pre (=elt?) (x))
  and (forall (x, y). post (=elt?) (x, y) (true) <=> x = y)
  and (forall (x, y). post (=elt?) (x, y) (false) <=> not (x = y))

(* There exists an order. *)
deferred predicate <elt : (elt * elt) ~> prop
where (forall x. not (x <elt x))
  and (forall (x, y). (not (x <elt y) and not (x = y)) => y <elt x)
  and (forall (x, y, z). (x <elt y and y <elt z) => x <elt z)
  and (forall (x, y). (x <elt y) or (y <elt x) or (x = y))

(* This order is decidable. *)
deferred val <elt? : (elt * elt) -> bool where
    (forall (x, y). (post (<elt?) (x, y) (true)  <=> x <elt y))
and (forall (x, y). (post (<elt?) (x, y) (false) <=> not (x <elt y)))
and (forall (x, y). pre (<elt?) (x, y))

(** The sets. *)

(* The type of an abstract set. *)
deferred type set : *

(* The membership predicate. *)
deferred predicate %in : (elt * set) ~> prop

(* The empty set. *)
deferred val empty : set
where forall x [x]. not (x %in empty)

deferred logic singleton : elt ~> set
where forall (x, y). x %in (singleton (y)) <=> (x = y)

deferred logic @ : (set * set) ~> set
where forall (s1, s2, x). x %in (s1 @ s2) <=> ((x %in s1) or (x %in s2))

deferred logic ^ : (set * set) ~> set
where forall (s1, s2, x). x %in (s1 ^ s2) <=> ((x %in s1) and (x %in s2))

deferred logic \\ : (set * set) ~> set
where forall (s1, s2, x). x %in (s1 \\ s2) <=> ((x %in s1) and not (x %in s2))

deferred predicate %C : (set * set) ~> prop
where forall (s1, s2). (s1 %C s2) <=> (forall x. (x %in s1) => (x %in s2))

lemma union_subset : forall (s1, s2). s1 %C (s1 @ s2)

lemma union_subset' : forall (s1, s2). s2 %C (s1 @ s2)

lemma flatten_union: forall (x, s1, s2, s3).
  x %in (s1 @ s2 @ s3) <=> (x %in s1 or x %in s2 or x %in s3)

lemma flatten_union': forall (x, s1, s2, s3).
  x %in (s1 @ (s2 @ s3)) <=> (x %in s1 or x %in s2 or x %in s3)

lemma transitive_subset: forall (s1, s2, s3).
  (s1 %C s2 and s2 %C s3) => s1 %C s3

(* Inferior bound of a set. *)
predicate inf (x, s) = forall y. y %in s => x <elt y

lemma: forall (x, s1, s2). (inf (x, s1) and s2 %C s1) => inf (x, s2)
lemma: forall (x, s1, s2). (inf (x, s1 @ s2)) <=> (inf (x, s1) and inf (x, s2))
lemma: forall (s, x, y). (inf (x, s) and y <elt x) => inf (y, s)
lemma: forall (x, y). (inf (x, singleton (y)) <=> x <elt y)
lemma: forall (x, s). (inf (x, s) => not (x %in s))

(* The minimum is unique. *)
lemma: forall (a, m1, m2).
      (m1 %in a and (forall y. (y %in a) => (m1 <elt y or m1 = y)) and
       m2 %in a and (forall y. (y %in a) => (m2 <elt y or m2 = y)))
      => m1 = m2

(* Superior bound of a set. *)
predicate sup (x, s) = forall y. y %in s => y <elt x

lemma: forall (x, s1, s2). (sup (x, s1) and s2 %C s1) => sup (x, s2)
lemma: forall (x, s1, s2). (sup (x, s1 @ s2)) <=> (sup (x, s1) and sup (x, s2))
lemma: forall (s, x, y). (sup (x, s) and x <elt y) => sup (y, s)
lemma: forall (x, y). (sup (x, singleton (y)) <=> y <elt x)
lemma: forall (x, s). (sup (x, s) => not (x %in s))

(* Extensional equality. *)
predicate === (s1, s2) = forall x. (x %in s1) <=> (x %in s2)

(* Extensionality axiom. *)
axiom: forall (s1, s2). (s1 === s2) <=> (s1 = s2)

deferred predicate disjoint : (set * set) ~> prop
where forall (s1, s2). disjoint (s1, s2) <=> ((forall y. y %in s1 => not (y %in s2)) and (forall y. y %in s2 => not (y %in s1)))

(** Trees. *)

(* Inductive type for tree. *)
type tree =
  | Empty : tree
  | Node  : (int * tree * elt * tree) -> tree

(* The abstraction of trees to sets. *)
logic elements (t) =
  | Empty -> empty
  | Node (h, l, x, r) -> elements (l) @ singleton (x) @ elements (r)

lemma root_in_elements: forall (h, l, x, r). x %in elements (Node (h, l, x, r))

lemma subl_in_elements: forall (h, l, x, r).
  elements (l) %C elements (Node (h, l, x, r))

lemma subr_in_elements: forall (h, l, x, r).
  elements (r) %C elements (Node (h, l, x, r))

lemma subl2_in_elements: forall (h, hl, ll, lx, lr, x, r).
    elements (ll) %C elements (Node (h, Node (hl, ll, lx, lr), x, r))
and elements (lr) %C elements (Node (h, Node (hl, ll, lx, lr), x, r))

lemma subr2_in_elements: forall (h, l, x, hr, rl, rx, rr).
  elements (rl) %C elements (Node (h, l, x, Node (hr, rl, rx, rr)))
and elements (rr) %C elements (Node (h, l, x, Node (hr, rl, rx, rr)))

(* A simple function: the emptyness test. *)
let is_empty t returns b where ((b = true) <=> (elements (t) = empty)) =
  match t with
    | Empty             -> true
    | Node (h, l, x, r) -> false
 end

(* Binary search trees. *)

inductive bst : tree ~> prop =
  | bst (Empty)
  | forall (h, l, x, r).
      bst (l) and bst (r) and sup (x, elements (l)) and inf (x, elements (r))
      => bst (Node (h, l, x, r))

predicate bst_node (l, x, r) =
   bst (l) and bst (r) and sup (x, elements (l)) and inf (x, elements (r))

lemma: forall (l, x, r). bst_node (l, x, r) =>
    ((disjoint (elements (l), elements (r)))
 and (disjoint (elements (l), singleton (x)))
 and (disjoint (elements (r), singleton (x))))

lemma: forall (lh, ll, lx, lr, x, rh, rl, rx, rr).
   bst_node (Node (lh, ll, lx, lr), x, Node (rh, rl, rx, rr)) =>
   (bst_node (ll, lx, lr) and bst_node (rl, rx, rr))

lemma: forall (l, x, h, rl, rx, rr).
   bst_node (l, x, Node (h, rl, rx, rr)) =>
   sup (rx, elements (l) @ singleton (x) @ elements (rl))

lemma: forall (h, ll, lx, lr, x, r).
   bst_node (Node (h, ll, lx, lr), x, r) =>
   inf (lx, elements (lr) @ singleton (x) @ elements (r))

lemma: forall (lh, ll, lx, lrh, lrl, lrx, lrr, x, r).
   bst_node (Node (lh, ll, lx, Node (lrh, lrl, lrx, lrr)), x, r) =>
   inf (lrx, elements (lrr) @ singleton (x) @ elements (r))

lemma: forall (h, l, x, rh, rlh, rll, rlx, rlr, rx, rr).
   bst_node (l, x, Node (rh, Node (rlh, rll, rlx, rlr), rx, rr)) =>
   sup (rlx, elements (l) @ singleton (x) @ elements (rll))

lemma: forall (h, l, x, r). bst (Node (h, l, x, r)) => (forall y. ((y %in elements (Node (h, l, x, r))) and y <elt x) <=> y %in elements (l))

lemma: forall (h, l, x, r). bst (Node (h, l, x, r)) => (forall y. ((y %in elements (Node (h, l, x, r))) and x <elt y) <=> y %in elements (r))

lemma: forall (h, l, x, r). bst_node (l, x, r) <=> bst (Node (h, l, x, r))

(* Membership decision procedure over binary search trees. *)

let rec mem_bst (t, x) where bst (t)
returns b where ((b = true) <=> (x %in elements (t))) =

    match t with

      | Empty ->
          false

      | Node (h, l, y, r) ->
          if (x =elt? y) then
            true
          else if (x <elt? y) then
            mem_bst (l, x)
          else
            mem_bst (r, x)

 end

(* Height. *)

logic height t =
  | Empty -> 0
  | Node (h, l, x, r) -> 1 + (max (height (l), height (r)))

lemma height_positive: forall t. height (t) >= 0

(* A tree is well-formed if the integer stored in its root node
   is its height. *)

inductive wf : tree ~> prop =
  | wf (Empty)
  | forall (h, l, x, r). h = height (Node (h, l, x, r)) and wf (l) and wf (r)
    => wf (Node (h, l, x, r))

(* We can extract the height of a well-formed tree in constant time. *)

let height? (t) where wf (t) returns h where (h = height (t)) =
  match t with
    | Empty -> 0
    | Node (h, l, x, r) -> h
  end

(* The definition of AVL trees. *)

predicate sim_depth (l, r) = dist (height (l), height (r), 2)

predicate almost_sim_depth (l, r) = dist (height (l), height (r), 3)

inductive avl : tree ~> prop =
  | avl (Empty)
  | forall (h, x, l, r).
      avl (l) and avl (r) and sim_depth (l, r)
    => avl (Node (h, l, x, r))

predicate avl_node (l, r) =
  avl (l) and avl (r) and sim_depth (l, r)

lemma: forall (h, l, x, r). avl_node (l, r) <=> avl (Node (h, l, x, r))

predicate almost_avl_node (l, r) =
  avl (l) and avl (r) and almost_sim_depth (l, r)

predicate wf_bst_avl (t) =
  wf (t) and bst (t) and avl (t)

(* Smart constructor of BST and AVL well-formed trees. *)

let mk_node (l, x, r)
where bst_node (l, x, r) and avl_node (l, r) and wf (l) and wf (r)
returns t
where (wf_bst_avl (t) and bst_node (l, x, r)
  and height (t) = max (height (l), height (r)) + 1
  and elements (t) = (elements (l) @ singleton (x) @ elements (r))
  and t = Node (max (height (l), height (r)) + 1, l, x, r)) =
  Node (max? (height? (l), height? (r)) +? 1, l, x, r)

(* Balancing. *)
axiom union_assoc : forall (s1, s2, s3). (s1 @ (s2 @ s3)) = ((s1 @ s2) @ s3)

let bal (l, x, r)
where wf (l) and wf (r) and almost_avl_node (l, r) and bst_node (l, x, r)
returns t
where (wf (t) and avl (t)
  and elements (t) = (elements (l) @ singleton (x) @ elements (r))
  and (avl_node (l, r) => height (t) = max (height (l), height (r)) + 1)
  and dist (1 + max (height (l), height (r)), height (t), 1)) =

  let hl = height? (l) in
  let hr = height? (r) in

    if (hl -? hr) =? 3 then let Node (h, ll, lx, lr) = l in

        if height? (lr) <=? height? (ll) then
	  mk_node (ll, lx, mk_node (lr, x, r))

        else let Node (h, lrl, lrx, lrr) = lr in
          mk_node (mk_node (ll, lx, lrl), lrx, mk_node (lrr, x, r))

    else if (hr -? hl) =? 3 then let Node (h, rl, rx, rr) = r in

        if height? (rl) <=? height? (rr) then
	  mk_node (mk_node (l, x, rl), rx, rr)

        else let Node (h, rll, rlx, rlr) = rl in
          mk_node (mk_node (l, x, rll), rlx, mk_node (rlr, rx, rr))

    else
      mk_node (l, x, r)
